{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "90710518",
   "metadata": {},
   "source": [
    "# L3: Advanced queries with Agents"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ee47ffb3",
   "metadata": {},
   "source": [
    "In this lesson, you are going to do more advanced querying."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0850beba",
   "metadata": {},
   "source": [
    "### Import required dependencies and load the API key"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0eb492b6-b561-4f95-882e-2d438692b21f",
   "metadata": {
    "height": 216
   },
   "outputs": [],
   "source": [
    "import * as mod from \"https://deno.land/std@0.213.0/dotenv/mod.ts\";\n",
    "import { \n",
    "    Document, \n",
    "    VectorStoreIndex, \n",
    "    SimpleDirectoryReader,\n",
    "    RouterQueryEngine,\n",
    "    OpenAIAgent,\n",
    "    QueryEngineTool,\n",
    "    FunctionTool,\n",
    "} from \"npm:llamaindex@0.1.18\"\n",
    "const keys = await mod.load({export:true}) // read API key from .env"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "57a8df97",
   "metadata": {},
   "source": [
    "Let's create 2 different data sources that answer different types of questions:\n",
    "- The first is the same example as in lesson 1, an essay by Dan Abramov about his career.\n",
    "- The second is the wikipedia page content about React.\n",
    "\n",
    "**Note**: To access the data used in this lesson, go to `File` and click on `Open` "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "838e4302-d735-48b6-bbec-a75dcec5cb15",
   "metadata": {
    "height": 80
   },
   "outputs": [],
   "source": [
    "const documents1 = await new SimpleDirectoryReader().loadData({directoryPath: \"./data\"})\n",
    "const index1 = await VectorStoreIndex.fromDocuments(documents1)\n",
    "const queryEngine1 = index1.asQueryEngine()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d4915eb5-98fe-4233-98cf-696d653a8552",
   "metadata": {
    "height": 63
   },
   "outputs": [],
   "source": [
    "let response1 = await queryEngine1.query({query: \"What did the author do in college?\"})\n",
    "console.log(response1.toString())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d2906d8e-c949-45a0-a496-f6b65f7cf846",
   "metadata": {
    "height": 80
   },
   "outputs": [],
   "source": [
    "const documents2 = await new SimpleDirectoryReader().loadData({directoryPath: \"./data2\"})\n",
    "const index2 = await VectorStoreIndex.fromDocuments(documents2)\n",
    "const queryEngine2 = index2.asQueryEngine()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "46868774-6909-4619-9643-b758e37105a2",
   "metadata": {
    "height": 63
   },
   "outputs": [],
   "source": [
    "let response2 = await queryEngine2.query({query:\"What is React?\"})\n",
    "console.log(response2.toString())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "89aed556",
   "metadata": {},
   "source": [
    "### Create a router query engine"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d04023ce-1551-4618-9dd3-6c157814004f",
   "metadata": {
    "height": 233
   },
   "outputs": [],
   "source": [
    "const queryEngine = await RouterQueryEngine.fromDefaults({\n",
    "  queryEngineTools: [\n",
    "    {\n",
    "      queryEngine: queryEngine1,\n",
    "      description: \"Useful for questions about Dan Abramov\",\n",
    "    },\n",
    "    {\n",
    "      queryEngine: queryEngine2,\n",
    "      description: \"Useful for questions about the React library\",\n",
    "    },\n",
    "  ],\n",
    "});"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "27e763cf",
   "metadata": {},
   "source": [
    "### Ask it a question"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "60917ee5-a48e-4507-a600-84afa5e8f225",
   "metadata": {
    "height": 63
   },
   "outputs": [],
   "source": [
    "let response3 = await queryEngine.query({query: \"What is React?\"})\n",
    "console.log(response3.toString())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc142f60-1989-4110-b0c8-1599939f557c",
   "metadata": {
    "height": 63
   },
   "outputs": [],
   "source": [
    "let response4 = await queryEngine.query({query: \"What did Dan Abramov do in college?\"})\n",
    "console.log(response4.toString())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4d4b4002",
   "metadata": {},
   "source": [
    "### Define a simple function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4a9ebe6c-3d8c-47d0-bdc3-48ccc379f050",
   "metadata": {
    "height": 63
   },
   "outputs": [],
   "source": [
    "function sumNumbers({a,b}) {\n",
    "  return a + b;\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e6cebe16",
   "metadata": {},
   "source": [
    "### Explain the function in a JSON structure"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "91010aa1-7288-4d9c-a55f-e4538345e414",
   "metadata": {
    "height": 250
   },
   "outputs": [],
   "source": [
    "const sumJSON = {\n",
    "  type: \"object\",\n",
    "  properties: {\n",
    "    a: {\n",
    "      type: \"number\",\n",
    "      description: \"The first number\",\n",
    "    },\n",
    "    b: {\n",
    "      type: \"number\",\n",
    "      description: \"The second number\",\n",
    "    },\n",
    "  },\n",
    "  required: [\"a\", \"b\"],\n",
    "};"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7b88cc1c",
   "metadata": {},
   "source": [
    "### Make the new function into a tool"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7a1f712a-d24e-4e38-bcf9-57a302f743c1",
   "metadata": {
    "height": 97
   },
   "outputs": [],
   "source": [
    "const sumFunctionTool = new FunctionTool(sumNumbers, {\n",
    "  name: \"sumNumbers\",\n",
    "  description: \"Use this function to sum two numbers\",\n",
    "  parameters: sumJSON,\n",
    "});"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fc946009",
   "metadata": {},
   "source": [
    "### Make the router query engine into a tool"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "abc6c491-43b8-4944-a383-46551a415790",
   "metadata": {
    "height": 148
   },
   "outputs": [],
   "source": [
    "const queryEngineTool = new QueryEngineTool({\n",
    "    queryEngine: queryEngine,\n",
    "    metadata: {\n",
    "        name: \"react_and_dan_abramov_engine\",\n",
    "        description: \"A tool that can answer questions about Dan Abramov and React\",\n",
    "    },\n",
    "});"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cf84c156",
   "metadata": {},
   "source": [
    "### Use the two tools"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c03fb86-d82e-4ab6-94ef-6b1407a4bb24",
   "metadata": {
    "height": 80
   },
   "outputs": [],
   "source": [
    "const agent = new OpenAIAgent({\n",
    "    tools: [queryEngineTool, sumFunctionTool],\n",
    "    verbose: true\n",
    "})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f1221ee1",
   "metadata": {},
   "source": [
    "### Let's ask it a question!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "70a427c6-3e0e-4216-8824-f229147e4209",
   "metadata": {
    "height": 63
   },
   "outputs": [],
   "source": [
    "let response5 = await agent.chat({message:\"What is React? Use a tool.\"})\n",
    "console.log(response5.toString())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c3460d6a-fbc0-4d93-8a79-529756fec52d",
   "metadata": {
    "height": 46
   },
   "outputs": [],
   "source": [
    "let response6 = await agent.chat({message:\"What is 501 + 5?\"})\n",
    "console.log(response6.toString())"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Deno",
   "language": "typescript",
   "name": "deno"
  },
  "language_info": {
   "file_extension": ".ts",
   "mimetype": "text/x.typescript",
   "name": "typescript",
   "nb_converter": "script",
   "pygments_lexer": "typescript",
   "version": "5.3.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
